import matplotlib.pyplot as plt
import torch
import numpy as np
from sklearn.datasets import load_diabetes  # 加载糖尿病数据集
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

# 加载糖尿病数据集
diabetes = load_diabetes()
data = diabetes.data
target = diabetes.target.reshape(-1, 1)

# 数据归一化
scaler_data = MinMaxScaler()
scaler_target = MinMaxScaler()
data = scaler_data.fit_transform(data)
target = scaler_target.fit_transform(target)

c = 7  # 基于你的代码，我们保持c=7不变
x = []
y = []
for i in range(len(data) - c):
    x_data = data[i:i + c]
    y_data = target[i + c]
    x.append(x_data)
    y.append(y_data)

x = torch.Tensor(x)
y = torch.Tensor(y)

# 数据划分
train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)

# 定义模型，注意糖尿病数据集每个样本有10个特征
model = torch.nn.Linear(in_features=7 * 10, out_features=1)  # 修改特征数量以适配糖尿病数据集

# 定义损失函数和优化器
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())

# 训练模型
loss_list = []
for epoch in range(2000):
    model.zero_grad()
    outputs = model(train_x.view(-1, 7 * 10))  # 修改输入形状以适配糖尿病数据集的特征数量
    loss = loss_fn(outputs, train_y)
    loss_list.append(loss.item())
    loss.backward()
    optimizer.step()
    if epoch % 10 == 0:
        print(f'Epoch {epoch + 1}, Loss: {loss.item()}')

# 预测测试集结果
test_outputs = model(test_x.view(-1, 7 * 10)).data.numpy()  # 修改输入形状以适配糖尿病数据集的特征数量

# 绘制结果
plt.plot(test_y.data.numpy(), c='r')
plt.plot(test_outputs, c='g')
plt.show()
